{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "ApdaLD4Qi30H"
      },
      "source": [
        "# Kuzu as Graph Memory"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "l7bi3i21i30I"
      },
      "source": [
        "## Prerequisites\n",
        "\n",
        "### Install Mem0 with Graph Memory support\n",
        "\n",
        "To use Mem0 with Graph Memory support, install it using pip:\n",
        "\n",
        "```bash\n",
        "pip install \"mem0ai[graph]\"\n",
        "```\n",
        "\n",
        "This command installs Mem0 along with the necessary dependencies for graph functionality.\n",
        "\n",
        "### Kuzu setup\n",
        "\n",
        "Kuzu comes embedded into the Python package that gets installed with the above command. There is no extra setup required.\n",
        "Just pick an empty directory where Kuzu should persist its database.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "DkeBdFEpi30I"
      },
      "source": [
        "## Configuration\n",
        "\n",
        "Do all the imports and configure OpenAI (enter your OpenAI API key):"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "id": "d99EfBpii30I"
      },
      "outputs": [],
      "source": [
        "from mem0 import Memory\n",
        "from openai import OpenAI\n",
        "\n",
        "import os\n",
        "\n",
        "os.environ[\"OPENAI_API_KEY\"] = \"\"\n",
        "openai_client = OpenAI()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "QTucZJjIi30J"
      },
      "source": [
        "Set up configuration to use the embedder model and Neo4j as a graph store:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 19,
      "metadata": {
        "id": "QSE0RFoSi30J"
      },
      "outputs": [],
      "source": [
        "config = {\n",
        "    \"embedder\": {\n",
        "        \"provider\": \"openai\",\n",
        "        \"config\": {\"model\": \"text-embedding-3-large\", \"embedding_dims\": 1536},\n",
        "    },\n",
        "    \"graph_store\": {\n",
        "        \"provider\": \"kuzu\",\n",
        "        \"config\": {\n",
        "            \"db\": \":memory:\",\n",
        "        },\n",
        "    },\n",
        "}\n",
        "memory = Memory.from_config(config_dict=config)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 20,
      "metadata": {},
      "outputs": [],
      "source": [
        "def print_added_memories(results):\n",
        "    print(\"::: Saved the following memories:\")\n",
        "    print(\" embeddings:\")\n",
        "    for r in results['results']:\n",
        "        print(\"    \",r)\n",
        "    print(\" relations:\")\n",
        "    for k,v in results['relations'].items():\n",
        "        print(\"    \",k)\n",
        "        for e in v:\n",
        "            print(\"      \",e)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "kr1fVMwEi30J"
      },
      "source": [
        "## Store memories\n",
        "\n",
        "Create memories:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 21,
      "metadata": {
        "id": "sEfogqp_i30J"
      },
      "outputs": [],
      "source": [
        "user = \"myuser\"\n",
        "\n",
        "messages = [\n",
        "    {\"role\": \"user\", \"content\": \"I'm planning to watch a movie tonight. Any recommendations?\"},\n",
        "    {\"role\": \"assistant\", \"content\": \"How about a thriller movies? They can be quite engaging.\"},\n",
        "    {\"role\": \"user\", \"content\": \"I'm not a big fan of thriller movies but I love sci-fi movies.\"},\n",
        "    {\"role\": \"assistant\", \"content\": \"Got it! I'll avoid thriller recommendations and suggest sci-fi movies in the future.\"}\n",
        "]"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "gtBHCyIgi30J"
      },
      "source": [
        "Store memories in Kuzu:"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 22,
      "metadata": {
        "id": "BMVGgZMFi30K"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "::: Saved the following memories:\n",
            " embeddings:\n",
            "     {'id': 'd3e63d11-5f84-4d08-94d8-402959f7b059', 'memory': 'Planning to watch a movie tonight', 'event': 'ADD'}\n",
            "     {'id': 'be561168-56df-4493-ab35-a5e2f0966274', 'memory': 'Not a big fan of thriller movies', 'event': 'ADD'}\n",
            "     {'id': '9bd3db2d-7233-4d82-a257-a5397cb78473', 'memory': 'Loves sci-fi movies', 'event': 'ADD'}\n",
            " relations:\n",
            "     deleted_entities\n",
            "     added_entities\n",
            "       [{'source': 'myuser', 'relationship': 'plans_to_watch', 'target': 'movie'}]\n",
            "       [{'source': 'movie', 'relationship': 'is_genre', 'target': 'thriller'}]\n",
            "       [{'source': 'movie', 'relationship': 'is_genre', 'target': 'sci-fi'}]\n",
            "       [{'source': 'myuser', 'relationship': 'has_preference', 'target': 'sci-fi'}]\n",
            "       [{'source': 'myuser', 'relationship': 'does_not_prefer', 'target': 'thriller'}]\n"
          ]
        }
      ],
      "source": [
        "results = memory.add(messages, user_id=user, metadata={\"category\": \"movie_recommendations\"})\n",
        "print_added_memories(results)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "id": "LBXW7Gv-i30K"
      },
      "source": [
        "## Search memories"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 23,
      "metadata": {
        "colab": {
          "base_uri": "https://localhost:8080/"
        },
        "id": "UHFDeQBEi30K",
        "outputId": "2c69de7d-a79a-48f6-e3c4-bd743067857c"
      },
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Loves sci-fi movies 0.31536642873409\n",
            "Planning to watch a movie tonight 0.0967911158879874\n",
            "Not a big fan of thriller movies 0.09468540071789472\n"
          ]
        }
      ],
      "source": [
        "for result in memory.search(\"what does alice love?\", user_id=user)[\"results\"]:\n",
        "    print(result[\"memory\"], result[\"score\"])"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Chatbot"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 24,
      "metadata": {},
      "outputs": [],
      "source": [
        "def chat_with_memories(message: str, user_id: str = user) -> str:\n",
        "    # Retrieve relevant memories\n",
        "    relevant_memories = memory.search(query=message, user_id=user_id, limit=3)\n",
        "    memories_str = \"\\n\".join(f\"- {entry['memory']}\" for entry in relevant_memories[\"results\"])\n",
        "    print(\"::: Using memories:\")\n",
        "    print(memories_str)\n",
        "\n",
        "    # Generate Assistant response\n",
        "    system_prompt = f\"You are a helpful AI. Answer the question based on query and memories.\\nUser Memories:\\n{memories_str}\"\n",
        "    messages = [{\"role\": \"system\", \"content\": system_prompt}, {\"role\": \"user\", \"content\": message}]\n",
        "    response = openai_client.chat.completions.create(model=\"gpt-4o-mini\", messages=messages)\n",
        "    assistant_response = response.choices[0].message.content\n",
        "\n",
        "    # Create new memories from the conversation\n",
        "    messages.append({\"role\": \"assistant\", \"content\": assistant_response})\n",
        "    results = memory.add(messages, user_id=user_id)\n",
        "    print_added_memories(results)\n",
        "\n",
        "    return assistant_response"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 25,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Chat with AI (type 'exit' to quit)\n",
            "::: Using memories:\n",
            "- Planning to watch a movie tonight\n",
            "- Not a big fan of thriller movies\n",
            "- Loves sci-fi movies\n",
            "::: Saved the following memories:\n",
            " embeddings:\n",
            " relations:\n",
            "     deleted_entities\n",
            "       []\n",
            "     added_entities\n",
            "       [{'source': 'myuser', 'relationship': 'loves', 'target': 'sci-fi'}]\n",
            "       [{'source': 'myuser', 'relationship': 'wants_to_avoid', 'target': 'thrillers'}]\n",
            "       [{'source': 'myuser', 'relationship': 'recommends', 'target': 'interstellar'}]\n",
            "       [{'source': 'myuser', 'relationship': 'recommends', 'target': 'the_martian'}]\n",
            "       [{'source': 'interstellar', 'relationship': 'is_a', 'target': 'sci-fi'}]\n",
            "       [{'source': 'the_martian', 'relationship': 'is_a', 'target': 'sci-fi'}]\n",
            "<<< AI: Since you love sci-fi movies and want to avoid thrillers, I recommend watching \"Interstellar\" if you haven't seen it yet. It's a visually stunning film that explores space travel, time, and love. Another great option is \"The Martian,\" which is more of a fun survival story set on Mars. Both films offer engaging stories and impressive visuals that are characteristic of the sci-fi genre!\n",
            "Goodbye!\n"
          ]
        }
      ],
      "source": [
        "print(\"Chat with AI (type 'exit' to quit)\")\n",
        "while True:\n",
        "    user_input = input(\">>> You: \").strip()\n",
        "    if user_input.lower() == 'exit':\n",
        "        print(\"Goodbye!\")\n",
        "        break\n",
        "    print(f\"<<< AI response:\\n{chat_with_memories(user_input)}\")"
      ]
    }
  ],
  "metadata": {
    "colab": {
      "provenance": []
    },
    "kernelspec": {
      "display_name": "mem0ai-sQeqgA1d-py3.12",
      "language": "python",
      "name": "python3"
    },
    "language_info": {
      "codemirror_mode": {
        "name": "ipython",
        "version": 3
      },
      "file_extension": ".py",
      "mimetype": "text/x-python",
      "name": "python",
      "nbconvert_exporter": "python",
      "pygments_lexer": "ipython3",
      "version": "3.12.10"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
